# coding = utf-8
import json
import os
from copy import deepcopy
from pathlib import PurePosixPath
import torch

import numpy as np
from PIL import Image

from models.experimental import attempt_load
from flask import Blueprint, request, send_from_directory

from models.unet import UNet
from utils.torch_utils import select_device

from models.yolopredict import yolo_predict_ret_img
from models.Unetpredict import unet_predict_img, lost_defect_detect, box_selectby_semantics, getContours
from configs.flask_id2name import id2name

lipstick_defect_detect = Blueprint('self', __name__)
show_choice = {}
showFlaws = None


with open('configs/flask_config.json', 'r', encoding='utf8') as fp:
    opt = json.load(fp)
    print('Flask Config:', opt)

device = select_device(opt['device'])
model = attempt_load(opt['weights'], map_location=device)


def update_json(opt):
    with open('configs/flask_config.json', 'w', encoding='utf8') as fp:
        fp.write(json.dumps(opt, ensure_ascii=False, indent=4))


def thresh_model(opt, show_choice):
    single_conf = True
    conf = opt['conf_thres']
    if type(conf) is not list:
        conf = 6 * [conf]

    if show_choice['face']:
        face_conf = opt['face_conf_thres']
        conf = [min(conf[i], face_conf[i]) for i in range(len(face_conf))]
        single_conf = False

    if show_choice['body']:
        body_conf = opt['body_conf_thres']
        conf = [min(conf[i], body_conf[i]) for i in range(len(body_conf))]
        single_conf = False

    if show_choice['bottom']:
        bottom_conf = opt['bottom_conf_thres']
        conf = [min(conf[i], bottom_conf[i]) for i in range(len(bottom_conf))]
        single_conf = False

    opt['thresh_model'] = conf
    if not single_conf:
        opt['conf_thres'] = min(conf)


@lipstick_defect_detect.route('/predict/', methods=["GET", 'POST', 'OPTIONS', 'PATCH'])
# 测试return请求信息
def predict():
    net = UNet(n_channels=3, n_classes=4)
    # image = request.files.get('image') # 通过文件形式获取image
    detect_json = request.get_json()
    # 校准参数
    config = detect_json.get('config')

    # image = open(detect_json.get('image'))
    img = Image.open(detect_json.get('image'))
    filename = PurePosixPath(detect_json.get('image')).name
    if PurePosixPath(filename).suffix == ".bmp":
        filename = PurePosixPath(filename).stem + ".jpg"
    img.convert("RGB").save(f'images/upload/{filename}')
    img = np.asarray(img)

    # 校准参数设置
    for k, v in config.items():
        opt[k] = v

    show_choice = {
        "face": config['face'],
        "body": config['body'],
        "bottom": config['bottom'],
    }
    showFlaws = config['showFlaws']
    thresh_model(opt, show_choice)

    img, results = yolo_predict_ret_img(opt, model, img, img_path=f'images/upload/{filename}', showFlaws=showFlaws)
    # ***************************Unet******************************
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net.to(device=device)
    net.load_state_dict(torch.load(opt["unet_weight"], map_location=device))

    img = Image.open(detect_json.get('image'))
    img = img.convert("RGB")
    # 调用模型
    mask = unet_predict_img(net=net, full_img=img, device=device, scale_factor=opt["unet_scale"],
                            out_threshold=opt["unet_mask_threshold"])

    background, bottom = mask[0, :, :] * 50, mask[1, :, :] * 50
    body, face = mask[2, :, :] * 50, mask[3, :, :] * 50
    body_contours, body_areas = getContours(body)
    face_contours, face_areas = getContours(face)
    bottom_contours, bottom_areas = getContours(bottom)
    # 检测断裂缺陷--新类型缺陷,与之前的并列
    results = lost_defect_detect(body_areas, face_areas, results)

    # 根据unet删选yolo的输出结果
    del_list = []
    new_id2name = dict(zip(id2name.values(), id2name.keys()))
    if results["results"] != []:
        for i, info in enumerate(results["results"]):
            temp_box = info["bbox"]
            # reverse key and values
            temp_cls = new_id2name[info["name"]]
            temp_conf = info["conf"]
            temp_box = [(temp_box[0] + temp_box[2]) // 2, (temp_box[1] + temp_box[3]) // 2]
            del_list = box_selectby_semantics(i, temp_box, body_contours[0], face_contours[0], bottom_contours[0],
                                              del_list, show_choice=show_choice, temp_cls=temp_cls, temp_conf=temp_conf, opt=opt)
    res = deepcopy(results["results"])
    for i in del_list:
        res[i] = 0
    while 0 in res:
        res.remove(0)
    results["results"] = res

    return results


@lipstick_defect_detect.route('/predict_file/', methods=["GET", 'POST', 'OPTIONS', 'PATCH'])
# 测试return请求信息
def predict_file():
    net = UNet(n_channels=3, n_classes=4)
    # 获取文件
    image = request.files.get('image')
    # 获取text
    config_text = request.form.get('config')
    config = json.loads(config_text).get('config')

    # image = open(detect_json.get('image'))
    img = Image.open(image)
    filename = PurePosixPath(image.filename).name
    if PurePosixPath(filename).suffix == ".bmp":
        filename = PurePosixPath(filename).stem + ".jpg"
    img.convert("RGB").save(f'upload/{filename}')
    img = np.asarray(img)

    # 校准参数设置
    for k, v in config.items():
        opt[k] = v

    show_choice = {
        "body": config['body'],
        "face": config['face'],
        "bottom": config['bottom'],
    }
    showFlaws = config['showFlaws']
    thresh_model(opt, show_choice)


    img, results = yolo_predict_ret_img(opt, model, img, img_path=f'upload/{filename}', showFlaws=showFlaws)
    # ***************************Unet******************************
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net.to(device=device)
    net.load_state_dict(torch.load(opt["unet_weight"], map_location=device))

    img = Image.open(image)
    img = img.convert("RGB")
    # 调用模型
    mask = unet_predict_img(net=net, full_img=img, device=device, scale_factor=opt["unet_scale"],
                            out_threshold=opt["unet_mask_threshold"])

    background, bottom = mask[0, :, :] * 50, mask[1, :, :] * 50
    body, face = mask[2, :, :] * 50, mask[3, :, :] * 50
    body_contours, body_areas = getContours(body)
    face_contours, face_areas = getContours(face)
    bottom_contours, bottom_areas = getContours(bottom)

    # 检测断裂缺陷--新类型缺陷,与之前的并列
    results = lost_defect_detect(body_areas, face_areas, results)

    # 根据unet删选yolo的输出结果
    del_list = []
    new_id2name = dict(zip(id2name.values(), id2name.keys()))
    if results["results"] != []:
        for i, info in enumerate(results["results"]):
            temp_box = info["bbox"]
            temp_box = [(temp_box[0] + temp_box[2]) // 2, (temp_box[1] + temp_box[3]) // 2]
            temp_cls = new_id2name[info["name"]]
            temp_conf = info["conf"]
            del_list = box_selectby_semantics(i, temp_box, body_contours[0], face_contours[0], bottom_contours[0],
                                              del_list, show_choice=show_choice, temp_cls=temp_cls, temp_conf=temp_conf, opt=opt)
    res = deepcopy(results["results"])
    for i in del_list:
        res[i] = 0
    while 0 in res:
        res.remove(0)
    results["results"] = res
    update_json(opt)
    return results





